package cn.jtcoding.web.framework.http;

import com.google.common.io.ByteStreams;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;

/**
 * POST方法的参数是存储在request.getInputStream中，只能读一次，不能多次读取。
 * 该类继承HttpServletRequestWrapper，先获取到请求参数，再将参数重新set到inputStream实现请求体多读。
 */
public class MultiReadHttpServletRequest extends HttpServletRequestWrapper {

	private HttpServletRequest originalRequest = null;
	private final byte[] bytes;
	private final WrappedServletInputStream wrappedServletInputStream;

	public MultiReadHttpServletRequest(HttpServletRequest request) throws IOException {
		super(request);
		this.originalRequest = request;
		// 读取输入流的请求参数，保存到bytes中
		this.bytes = ByteStreams.toByteArray(request.getInputStream());
		ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(this.bytes);
		this.wrappedServletInputStream = new WrappedServletInputStream(byteArrayInputStream);
	}

	/**
	 * 把参数重新写进请求里
	 */
	public void reWriteInputStream() {
		wrappedServletInputStream.setStream(new ByteArrayInputStream(bytes != null ? bytes : new byte[0]));
	}
	
	public void reWriteInputStream(byte[] bytes) {
		wrappedServletInputStream.setStream(new ByteArrayInputStream(bytes != null ? bytes : new byte[0]));
	}

	@Override
	public ServletInputStream getInputStream() throws IOException {
		return wrappedServletInputStream;
	}

	@Override
	public BufferedReader getReader() throws IOException {
		return new BufferedReader(new InputStreamReader(wrappedServletInputStream));
	}

	/**
	 * 获取POST/PUT参数
	 */
	public String getRequestBody() throws IOException {
		return new String(bytes, this.getCharacterEncoding());
	}

	// 清洗参数，防止xss注入
	public String[] getXssEncodeParameterValues(String parameter) {
		String[] values = super.getParameterValues(parameter);
		if (values == null) {
			return null;
		}
		int count = values.length;
		String[] encodedValues = new String[count];
		for (int i = 0; i < count; i++) {
			encodedValues[i] = xssEncode(values[i]);
		}
		return encodedValues;
	}

	/**
	 * 将Parameter的参数名和参数值都做xss过滤。<br/>
	 * 如果需要获得原始的值，则通过super.getParameterValues(name)来获取<br/>
	 * 
	 * @param name
	 * @return
	 */
	public String getXssEncodeParameter(String name) {
		String value = super.getParameter(xssEncode(name));
		if (value != null) {
			value = xssEncode(value);
		}
		return value;
	}

	/**
	 * 将Header的参数名和参数值都做xss过滤。<br/>
	 * 如果需要获得原始的值，则通过super.getHeaders(name)来获取<br/>
	 */
	public String getXssEncodeHeader(String name) {
		String value = super.getHeader(xssEncode(name));
		if (value != null) {
			value = xssEncode(value);
		}
		return value;
	}

	/**
	 * 获取最原始的HttpServletRequest
	 *
	 * @return
	 */
	public HttpServletRequest getOriginalRequest() {
		return originalRequest;
	}

	/**
	 * 将容易引起xss漏洞的半角字符直接替换成全角字符
	 * 
	 * @param str
	 * @return
	 */
	private static String xssEncode(String str) {
		if (str == null || "".equals(str)) {
			return str;
		}
		StringBuilder sb = new StringBuilder(str.length() + 16);
		for (int i = 0; i < str.length(); i++) {
			char c = str.charAt(i);
			switch (c) {
			case '>':
				sb.append('＞');// 全角大于号
				break;
			case '<':
				sb.append('＜');// 全角小于号
				break;
			case '\'':
				sb.append('‘');// 全角单引号
				break;
			case '\"':
				sb.append('“');// 全角双引号
				break;
			case '&':
				sb.append('＆');// 全角
				break;
			case '\\':
				sb.append('＼');// 全角斜线
				break;
			case '#':
				sb.append('＃');// 全角井号
				break;
			default:
				sb.append(c);
				break;
			}
		}
		return sb.toString();
	}

	private static class WrappedServletInputStream extends ServletInputStream {
		private InputStream stream;

		public WrappedServletInputStream(InputStream stream) {
			this.stream = stream;
		}

		@Override
		public int read() throws IOException {
			return stream.read();
		}

		public void setStream(InputStream stream) {
			this.stream = stream;
		}

		@Override
		public boolean isFinished() {
			try {
				return stream.available() == 0;
			} catch (IOException e) {
			}
			return true;
		}

		@Override
		public boolean isReady() {
			return true;
		}

		@Override
		public void setReadListener(ReadListener readListener) {

		}
	}
}
